import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

df = pd.read_excel('./associationRules/餐厅数据.xlsx')
print(df.head())

# 转换菜品格式
transactions = df['菜品'].str.split(',').tolist()

te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

# 使用apriori进行关联分析
frequent_itemsets = apriori(df_encoded, min_support=0.1, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 选择项集
print(frequent_itemsets[frequent_itemsets.itemsets.apply(lambda x: len(x) == 2)])

# 生成关联规则
rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=0.1)
effective = [rule for rule in rules if (rule['lift'] > 1 and rule['confidence'] > 0.1)]
effective.sort_values(by=['lift', 'confidence'], ascending=False)

print(effective[['antecedents', 'consequents', 'support', 'confidence', 'lift']])

import matplotlib.pyplot as plt
import networkx as nx

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

G = nx.DiGraph()

for _, row in effective.iterrows():
    G.add_edge(
        ', '.join(list(row['antecedents'])),
        ', '.join(list(row['consequents'])), 
        weight=row['lift']
    )

pos = nx.spring_layout(G)

nx.draw_networkx_edges(
    G,
    pos,
    edge_color=[v[v.index('weight')] for u, v in G.edges()],
    width=2.0,
    edge_cmap=plt.cm.Blues
)

plt.show()